import os
from scipy.stats import norm
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt



def rv_value(k,n,p,m):
    x= (k-m*n*p)/np.sqrt(m*n*p*(1-p))
    return x

def rv_value_difference(l,n,p,m):
    x= (l-(m-1)*n*p)/np.sqrt((m+1)*n*p*(1-p))
    return x

generalization_results = 'generalization_results/'
a = os.path.exists(generalization_results)
if not a:
    os.makedirs(generalization_results)


n=100
p=0.5
m=1
k=np.array(list(range(500)))



x1= rv_value(k,n,p,m)
p1 = 1-norm.cdf(x1, loc=0, scale=1)

x2= rv_value(k,n,p,m+5)
p2 = 1-norm.cdf(x2, loc=0, scale=1)

x3= rv_value_difference(k,n,p,m+5)
p3 = 1-norm.cdf(x3, loc=0, scale=1)


fig1 = plt.figure()
plt.plot(k,p1,label='Single-agent: '+r"$P\{Dim_{0}\geq k\}$",
         marker='v', markersize=5,linestyle = '-')
plt.plot(k,p2,label='Multi-agent: '+r"$P\{\sum_{i=1}^{m} Dim_{i}\geq k\}$",
         marker='x', markersize=5,linestyle = '-.')

plt.xlabel('Frequency: '+r"$k$")
plt.ylabel('Probability')
plt.legend(loc="upper right")
file_1 = 'gene_1.png'
plt.savefig(generalization_results+file_1)

k=np.array(list(range(500)))
m1= [5,8,10,20]
marklist = ['x','^','+','v']
llist = ['-','-','-','-']
x_list = []
p_list = []
for m_i in m1:
    x_i= rv_value_difference(k,n,p,m+m_i)
    p_i = 1-norm.cdf(x_i, loc=0, scale=1)
    x_list.append(x_i)
    p_list.append(p_i)

fig2 = plt.figure()
ax = fig2.add_subplot(111)
for i,id in enumerate(m1):
    ax.plot(k,p_list[i],label='Agent_num:{}'.format(id+1),
            marker=marklist[i], markersize=5,linestyle = llist[i])
ax.text(2,0.5,r"$P\{\sum_{i=1}^{m} Dim_{i}-Dim_{0}\geq l\}$")
ax.set_xlabel('Exceedance: '+r"$l$")
ax.set_ylabel('Probability')
plt.legend()

file_2 = 'difference_2.png'
plt.savefig(generalization_results+file_2)

# fig3 = plt.figure()
# plt.plot(k,p4,label='Difference')
# plt.xlabel('Exceedance: '+r"$l$")
# plt.ylabel('Probability')
# plt.legend()
# file_3 = 'difference_2.png'
# plt.savefig(generalization_results+file_3)
# # plt.show()